/**
 * @file add_custom.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "leakyrelugrad_custom_tiling.h"
#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2;

constexpr uint32_t CUSTOM_BFLOAT16 = 0;
constexpr uint32_t CUSTOM_FLOAT16 = 1;
constexpr uint32_t CUSTOM_FLOAT32 = 2;
constexpr uint32_t CUSTOM_INT8 = 3;
constexpr uint32_t CUSTOM_INT16 = 4;
constexpr uint32_t CUSTOM_INT32 = 5;

constexpr uint32_t LAST_TWO_TILE = 2;



template <typename dataType> 
class KernalLeakyReluGrad {
public:
    __aicore__ inline KernalLeakyReluGrad() {}
    __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR dx, float negativeSlope, LeakyReluGradCustomTilingData tiling)
    {
        this->negativeSlope = negativeSlope;
        if (tiling.isEvenCore) {
            this->blockLength = tiling.blockLength;
            this->tileNum = tiling.tileNum;
            this->tileLength = tiling.tileLength / BUFFER_NUM;
            this->lastTileLength = tiling.lastTileLength;

            dyGm.SetGlobalBuffer((__gm__ dataType *)dy + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
            xGm.SetGlobalBuffer((__gm__ dataType *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
            dxGm.SetGlobalBuffer((__gm__ dataType *)dx + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        } else {
            if (AscendC::GetBlockIdx() < tiling.formerNum) {
                this->tileNum = tiling.formerTileNum;
                this->tileLength = tiling.formerTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.formerLastTileLength;

                dyGm.SetGlobalBuffer((__gm__ dataType *)dy + tiling.formerLength * AscendC::GetBlockIdx(), tiling.formerLength);
                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * AscendC::GetBlockIdx(), tiling.formerLength);
                dxGm.SetGlobalBuffer((__gm__ dataType *)dx + tiling.formerLength * AscendC::GetBlockIdx(), tiling.formerLength);
            } else {
                this->tileNum = tiling.tailTileNum;
                this->tileLength = tiling.tailTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.tailLastTileLength;

                dyGm.SetGlobalBuffer((__gm__ dataType *)dy + tiling.formerLength * tiling.formerNum + 
                    tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum), tiling.tailLength);
                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * tiling.formerNum + 
                    tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum), tiling.tailLength);
                dxGm.SetGlobalBuffer((__gm__ dataType *)dx + tiling.formerLength * tiling.formerNum + 
                    tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum), tiling.tailLength);
            }
        }
        pipe.InitBuffer(inQueueDy, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(tmpBuffer, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(outQueueDx, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(tmp1Buffer, this->tileLength * sizeof(dataType));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<dataType> dyLocal = inQueueDy.AllocTensor<dataType>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1))) {
            AscendC::DataCopy(dyLocal, dyGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength],
                this->tileLength);
        } else {
            AscendC::DataCopy(dyLocal, dyGm[progress * this->tileLength], this->tileLength);
        }
        inQueueDy.EnQue(dyLocal);
        AscendC::LocalTensor<dataType> xLocal = inQueueX.AllocTensor<dataType>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1))) {
            AscendC::DataCopy(xLocal, xGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength],
                this->tileLength);
        } else {
            AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        }
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<dataType> dyLocal = inQueueDy.DeQue<dataType>();
        AscendC::LocalTensor<dataType> xLocal = inQueueX.DeQue<dataType>();
        AscendC::LocalTensor<dataType> dxLocal = outQueueDx.AllocTensor<dataType>();
        AscendC::LocalTensor<uint8_t> temp = tmpBuffer.Get<uint8_t>();
        AscendC::LocalTensor<dataType> negativePart = tmp1Buffer.Get<dataType>();

        AscendC::CompareScalar(temp, xLocal, static_cast<dataType>(0.0f), AscendC::CMPMODE::GT,this->tileLength);
        AscendC::DataCopy(negativePart, dyLocal, this->tileLength);
        AscendC::Muls(negativePart, negativePart, static_cast<dataType>(this->negativeSlope), this->tileLength);
        AscendC::Select(dxLocal, temp, dyLocal, negativePart, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, this->tileLength);
        
        outQueueDx.EnQue<dataType>(dxLocal);

        inQueueDy.FreeTensor(dyLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<dataType> dxLocal = outQueueDx.DeQue<dataType>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1))) {
            AscendC::DataCopy(dxGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength], dxLocal,
                this->tileLength);
        } else {
            AscendC::DataCopy(dxGm[progress * this->tileLength], dxLocal, this->tileLength);
        }
        outQueueDx.FreeTensor(dxLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueDy;  
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;   
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueDx; 
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer;            
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmp1Buffer;

    AscendC::GlobalTensor<dataType> dyGm;  
    AscendC::GlobalTensor<dataType> xGm;   
    AscendC::GlobalTensor<dataType> dxGm;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;
    float negativeSlope;

};
        
extern "C" __global__ __aicore__ void leakyrelugrad_custom(GM_ADDR dy, GM_ADDR x, GM_ADDR dx, float negativeSlope, LeakyReluGradCustomTilingData tiling)
{

    if (tiling.dataType == CUSTOM_FLOAT32) {
        KernalLeakyReluGrad<float> op;
        op.Init(dy, x, dx, negativeSlope, tiling);
        op.Process();
    }else if (tiling.dataType == CUSTOM_FLOAT16) {
        KernalLeakyReluGrad<half> op;
        op.Init(dy, x, dx, negativeSlope, tiling);
        op.Process();
   } 

}

